package ai

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

type DkReply struct {
	ID      string `json:"id"`
	Choices []struct {
		Delta struct {
			Content          string `json:"content"`
			ReasoningContent string `json:"reasoning_content"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason"`
		Index        int    `json:"index"`
		Logprobs     string `json:"logprobs"`
	} `json:"choices"`
	Created           int    `json:"created"`
	Model             string `json:"model"`
	SystemFingerprint string `json:"system_fingerprint"`
	Object            string `json:"object"`
}

type ChatReply struct {
	ID      string `json:"id"`
	Status  uint   `json:"status"`  //0未开始,1进行中;2已结束;
	Type    uint   `json:"type"`    //0无;1思考;2结果
	Content string `json:"content"` //内容
	Err     error  `json:",omitempty"`
}

type Chat struct {
	Url string `json:"url"` //api
	Key string `json:"key"` //密钥
}

func NewChat(url string, key string) *Chat {
	return &Chat{
		Url: url,
		Key: key,
	}
}
func (c Chat) newRequest(question string) (res *http.Response, e error) {
	payload := fmt.Sprintf(`{
		"model":"deepseek-r1",
		"messages":[
			{
				"role": "user", 
            	"content": "%s"
			}
		],
		"stream": true,
		"stream_options": {
			"include_usage": true
		}
	}`, question)

	req, err := http.NewRequest(http.MethodPost, c.Url, strings.NewReader(payload))
	if err != nil {
		e = err
		return
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", c.Key))

	return http.DefaultClient.Do(req)
}

func (c Chat) Question(ctx context.Context, question string, stream chan<- *ChatReply) {

	res, err := c.newRequest(question)
	if err != nil {
		stream <- &ChatReply{
			Err: err,
		}
		return
	}

	defer res.Body.Close()

	//处理流式响应
	scanner := bufio.NewReader(res.Body)

	for {
		select {
		case <-ctx.Done():
			break
		default:
			line, err := scanner.ReadString('\n')
			if err != nil {
				if err == io.EOF {
					break
				} else {
					stream <- &ChatReply{
						Err: err,
					}
					return
				}
			}

			line = strings.TrimSpace(line)
			if line == "" {
				break
			}

			if line == "data: [DONE]" {
				stream <- &ChatReply{
					Status: 2,
				}
				return
			}

			if strings.HasPrefix(line, "data: ") {
				line = strings.TrimPrefix(line, "data: ")
			}

			var response DkReply
			if err := json.Unmarshal([]byte(line), &response); err != nil {
				continue
			}

			if len(response.Choices) > 0 {
				var (
					t uint
					c string
				)
				if response.Choices[0].Delta.ReasoningContent != "" {

					t = 1
					c = response.Choices[0].Delta.ReasoningContent
				} else if response.Choices[0].Delta.Content != "" {

					t = 2
					c = response.Choices[0].Delta.Content
				}

				if c != "" {
					stream <- &ChatReply{
						ID:      response.ID,
						Status:  1,
						Type:    t,
						Content: c,
					}
				}
			}
		}
	}

}
